import torch
import torchvision
import torchvision.transforms as transforms
from collections import Counter

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载MNIST训练集和测试集
trainset = torchvision.datasets.MNIST(root='./data/mnist', train=True,
                                      download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=len(trainset),
                                          shuffle=False)

testset = torchvision.datasets.MNIST(root='./data/mnist', train=False,
                                     download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=len(testset),
                                         shuffle=False)

# 获取训练数据和标签
train_images, train_labels = next(iter(trainloader))
train_images = train_images.view(train_images.size(0), -1)  # 展平图像

# 获取测试数据和标签
test_images, test_labels = next(iter(testloader))
test_images = test_images.view(test_images.size(0), -1)  # 展平图像

# 定义KNN函数
def knn_predict(test_images, train_images, train_labels, k):
    predictions = []
    for test_image in test_images:
        # 计算测试样本与所有训练样本的欧氏距离
        distances = torch.cdist(test_image.unsqueeze(0), train_images).squeeze()
        # 获取距离最近的K个样本的索引
        _, indices = torch.topk(distances, k, largest=False)
        # 获取这K个样本的标签
        k_nearest_labels = train_labels[indices]
        # 统计每个标签出现的次数
        label_counts = Counter(k_nearest_labels.tolist())
        # 选择出现次数最多的标签作为预测结果
        prediction = label_counts.most_common(1)[0][0]
        predictions.append(prediction)
    return torch.tensor(predictions)

# 设置K值
import sys

k = int(sys.argv[1]) if len(sys.argv) > 1  else 5
print(f"k = {k}")

# 进行预测
predictions = knn_predict(test_images, train_images, train_labels, k)

# 计算准确率
correct = (predictions == test_labels).sum().item()
total = test_labels.size(0)
accuracy = correct / total
print(f'KNN分类准确率 (k={k}): {accuracy * 100:.2f}%')
